package problems.practice;

import java.util.HashSet;

/**
 * 52. N皇后II
 * <p>https://leetcode.cn/problems/n-queens-ii/</p>
 *
 * @author habitplus
 * @since 12:07, 2022/10/1
 */
public class T52 {
    // 解题思路：回溯法
    // HashSet
    public int totalNQueens1(int n) {
        HashSet<Integer> columns = new HashSet<Integer>();
        HashSet<Integer> diagonals1 = new HashSet<Integer>();
        HashSet<Integer> diagonals2 = new HashSet<Integer>();
        return backtrack(n, 0, columns, diagonals1, diagonals2);
    }

    public int backtrack(int n, int row, HashSet<Integer> columns, HashSet<Integer> diagonals1, HashSet<Integer> diagonals2) {
        if (row == n) {
            return 1;
        } else {
            int count = 0;
            int dia1, dia2;
            for (int i = 0; i < n; i++) {
                if (columns.contains(i)) continue;

                dia1 = row - i;
                if (diagonals1.contains(dia1)) continue;

                dia2 = row + i;
                if (diagonals2.contains(dia2)) continue;

                columns.add(i);
                diagonals1.add(dia1);
                diagonals2.add(dia2);
                count += backtrack(n, row + 1, columns, diagonals1, diagonals2);
                // 还原现场
                diagonals2.remove(dia2);
                diagonals1.remove(dia1);
                columns.remove(i);
            }
            return count;
        }
    }


    // 位运算
    public int totalNQueens(int n) {
        return dfs(n, 0, 0, 0, 0);
    }

    private int dfs(int n, int row, int col, int dia1, int dia2) {
        if (row >= n) return 1;
        // bit = 1 表示已经放置了皇后
        // col: 已放置皇后的列
        // row: 当前是放置第几个皇后，从 0 开始
        // dia1: 正斜线上已放置的皇后
        // dia2: 反斜线上已放置的皇后

        int count = 0;
        int bits = ((1 << n) - 1) & (~(col | dia1 | dia2));
        int pos;

        while (bits != 0) {
            pos = bits & (-bits);
            bits = bits & (bits - 1);
            count += dfs(n, row + 1, col | pos, (dia1 | pos) << 1, (dia2 | pos) >> 1);
        }

        return count;
    }
}
